import torch
import torch.nn.functional as F

from torch import nn

class LocalContrastiveLossIn(nn.Module):
    def __init__(self, 
                 T: int = 5,
                 std: float = 0.0,
                 temperature: float = 0.1,
                 alpha: float = 0.5
                 ) -> None:
        super().__init__()
        self.temperature = temperature
        self.std = std
        self.T = T
        self.alpha = alpha
        self.pos_prob = None

    def _init_pos_prob(self, device) -> None:
        if self.std == 0.0:
            pos_prob = torch.eye(self.T, device=device).float()
        else:
            Ts = torch.linspace(0., 1., self.T, device=device)
            dist = Ts[:, None] - Ts[None, :]

            prob_scores = torch.exp(-0.5 * (dist / self.std) ** 2)
            pos_prob = prob_scores / prob_scores.sum(dim=-1, keepdim=True)
        
        self.pos_prob = pos_prob

    def forward(self, z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor:
        """
        [b, T, d]
        """
        device = z1.device
        if self.pos_prob is None or self.pos_prob.device != device:
            self._init_pos_prob(device)

        z1_norm = F.normalize(z1, p=2, dim=-1)
        z2_norm = F.normalize(z2, p=2, dim=-1)

        scores_12 = torch.einsum('ntd,nkd->ntk', z1_norm, z2_norm) / self.temperature

        pos_scores_12 = torch.logsumexp(
            torch.log(self.pos_prob[None, :, :]) + scores_12, dim=-1
        )
        loss_12 = -pos_scores_12 + torch.logsumexp(scores_12, dim=-1)

        scores_21 = scores_12.transpose(1, 2)

        pos_scores_21 = torch.logsumexp(
            torch.log(self.pos_prob[None, :, :]) + scores_21, dim=-1
        )
        loss_21 = -pos_scores_21 + torch.logsumexp(scores_21, dim=-1)

        loss_12 = loss_12.mean()
        loss_21 = loss_21.mean()

        return self.alpha * loss_12 + (1.0 - self.alpha) * loss_21


class LocalContrastiveLossOut(nn.Module):
    def __init__(
        self,
        T: int = 10,
        temperature: float = 0.1,
        std: float = 0.0,
        alpha: float = 0.5
        ):
        
        super().__init__()       
        self.temperature = temperature
        self.alpha = alpha
        self.std = std

        if std == 0.0:
            pos_prob = torch.eye(T).float()
        else:
            Ts = torch.linspace(0., 1., T)
            # T, T
            dist = Ts[:, None] - Ts[None, :]

            prob_scores = torch.exp(-0.5 * (dist / std) ** 2)
            pos_prob = prob_scores / prob_scores.sum(dim=-1, keepdim=True)

        self.register_buffer('pos_prob', pos_prob)
    
    def forward(self, z1, z2):
        z1 = F.normalize(z1, p=2, dim=-1) # N, T, D
        z2 = F.normalize(z2, p=2, dim=-1) # N, T, D

        # N, T, T
        scores_12 = torch.einsum('ntd,nkd->ntk', z1, z2) / self.temperature

        # N, T
        pos_scores_12 = (self.pos_prob[None, :, :] * scores_12).sum(-1)
        # N, T
        loss_12 = -pos_scores_12 + torch.logsumexp(scores_12, dim=-1)

        scores_21 = scores_12.transpose(1, 2)
        pos_scores_21 = (self.pos_prob[None, :, :] * scores_21).sum(-1)
        loss_21 = -pos_scores_21 + torch.logsumexp(scores_21, dim=-1)

        loss_12 = loss_12.mean()
        loss_21 = loss_21.mean()

        return self.alpha * loss_12 + (1. - self.alpha) * loss_21


class LocalLossCosineSimilarity(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, z1, z2):
        """
        z1: Tensor of shape [N, T, D]
        z2: Tensor of shape [N, T, D]
        """
        z1_normalized = F.normalize(z1, dim=-1)
        z2_normalized = F.normalize(z2, dim=-1)

        cos_sim = torch.einsum('ntd,ntd->nt', z1_normalized, z2_normalized)  # Shape: [N, T]

        loss = -cos_sim.mean()

        return loss